import numpy as np
from enum import Enum
from ._base import EquatorialWave
from ..physics.constant import PI, EARTH_GRAVITY, ROSSBY_PARAMETER_ON_EQUATOR
from pint import Quantity
from .. import unit
from ..physics.unit import _handle_quantity


def dimensionless_wave_omega(wave: EquatorialWave, k: np.ndarray, n: int=1,
                             reverse_k: bool=True,):

    if reverse_k:
        k = -k

    if wave == EquatorialWave.MixedRossbyGravity:
        return np.sqrt(k**2/4 + 1) - k/2
    elif wave == EquatorialWave.Kelvin:
        return -k

    k[k == 0] = 1e-10
    k = np.asarray(k, dtype=np.complex64)
    # 把 k 的 0 转为非常小的数
    p = -(k**2 + 2*n + 1)
    q = k
    o = (-1+3**(1/2)*1j)/2

    if wave == EquatorialWave.EastwardInertioGravity:
        omega = (-q/2+((q/2)**2+(p/3)**3)**(1/2))**(1/3) +\
                (-q/2-((q/2)**2+(p/3)**3)**(1/2))**(1/3)
    elif wave == EquatorialWave.WestwardInertioGravity:
        omega = o * (-q/2+((q/2)**2+(p/3)**3)**(1/2))**(1/3) +\
                o**2 * (-q/2-((q/2)**2+(p/3)**3)**(1/2))**(1/3)
    elif wave == EquatorialWave.EquatorialRossby:
        omega = o**2 * (-q/2+((q/2)**2+(p/3)**3)**(1/2))**(1/3) +\
                o * (-q/2-((q/2)**2+(p/3)**3)**(1/2))**(1/3)

    return np.real(omega)


def dimensionless_n(wave: EquatorialWave, k: np.ndarray, omega: np.ndarray,
                    reverse_k: bool=True):

    if reverse_k:
        k = -k

    if k.ndim == 1 and omega.ndim == 1:
        K, O = np.meshgrid(k, omega)

    n_values = (O**2 - K**2 + K/O - 1)/2
    if wave == EquatorialWave.EastwardInertioGravity:
        wave_omega = dimensionless_wave_omega(wave, k, 1, reverse_k=False)
        n_values[O < wave_omega-0.5] = np.nan
    elif wave == EquatorialWave.EquatorialRossby:
        wave_omega = dimensionless_wave_omega(wave, k, 1, reverse_k=False)
        n_values[O > wave_omega+0.2] = np.nan
    else:
        raise ValueError("only EastwardInertioGravity and EquatorialRossby " \
        "waves are supported. Use dimensionless_wave_omega instead.")

    return n_values



def wave_omega(wave: EquatorialWave, k: Quantity | np.ndarray | float, n: int=1,
          inverse_k: bool=True, H: Quantity|float=25, g: Quantity=None,
          beta: Quantity=None) -> Quantity:
    """根据波数计算赤道波动的角频率，单位为 rad/s。公式来自 Matsuno (1966)。
    如果波数 k 无单位，则认为是 '1/m'，如果想使用纬向波数，可以使用单位
    'zonal_wavenumber' 来定义。

    返回的频率与 wave_frequency 函数的频率相差 2 * PI 倍。

    >>> from moisten import unit
    >>> from moisten.equatorialWaves import EquatorialWave, wave_omgea
    >>> k = np.linspace(-20, 20, 100) * unit('zonal_wavenumber')
    >>> omega = wave_omgea(EquatorialWave.MixedRossbyGravity, k)

    Parameters
    ----------
    wave : EquatorialWave
        波动类型

        The type of wave. See the `EquatorialWave` enum for options.
    k : np.ndarray | float
        波数, 如果没有单位，则认为是 '1/m'。这里的波数是指单位距离内的波动周期次数，
        而不是相位差。

        The wavenumber, which is the number of waves per unit distance.
        If no unit is given, it is assumed to be in '1/m'.
    n : int, optional
        频散关系中的 n。当波动类型为 MRG 与 Kelvin 波时，该值无效。

        The value of n in the dispersion relation. This is ignored for MRG and Kelvin waves.
    inverse_k : bool, optional
        是否反转波数的正负方向，默认为 True。当为 True 时，输入的波数 > 0 表示波动
        向东传播，反之为 False 时，输入的波数 > 0 表示波动向西传播。

        Whether to invert the sign of the wavenumber. Default is True.
        When True, a positive wavenumber indicates eastward propagation,
        and when False, a positive wavenumber indicates westward propagation.
    H : Quantity | float, optional
        相当深度，默认单位为米。

        Equivalent depth, default unit is meters.
    g : Quantity, optional
        重力加速度，如果不提供，则使用默认的地球重力加速度。

        Gravitational acceleration. If not provided, the earth's gravitational
        acceleration is used.
    beta : Quantity, optional
        赤道上的罗斯贝参数，如果不提供，则使用默认的地球上的罗斯贝参数。

        Rossby parameter at the equator. If not provided, the default value
        for the earth is used.

    Returns
    -------
    Quantity
        波动的角频率，单位为 rad/s。注意这是角频率，是指单位时间内的相位变化。

        The angular frequency of the wave, in radians per second.
    """

    # 把 k 转为波数（个/米）
    k = _handle_quantity(k, '1/m')
    # 转为相位/米的波数
    k = k * 2 * PI

    # 处理常数
    if g is None:
        g = EARTH_GRAVITY
    if beta is None:
        beta = ROSSBY_PARAMETER_ON_EQUATOR
    H = _handle_quantity(H, 'm')

    if inverse_k:
        k *= -1

    # 重力波速度
    c = np.sqrt(g*H)
    sqrt2 = np.sqrt(beta/(2*c))

    # 转换
    k = k * np.sqrt(c / beta)
    k = k.astype(np.complex64)
    k[k == 0] = 1e-10

    p = -(k**2 + 2*n + 1)
    q = k
    o = (-1+3**(1/2)*1j)/2

    match wave:
        case EquatorialWave.EastwardInertioGravity:
            omega = (-q/2+((q/2)**2+(p/3)**3)**(1/2))**(1/3) +\
                    (-q/2-((q/2)**2+(p/3)**3)**(1/2))**(1/3)
        case EquatorialWave.WestwardInertioGravity:
            omega = o * (-q/2+((q/2)**2+(p/3)**3)**(1/2))**(1/3) +\
                    o**2 * (-q/2-((q/2)**2+(p/3)**3)**(1/2))**(1/3)
        case EquatorialWave.EquatorialRossby:
            omega = o**2 * (-q/2+((q/2)**2+(p/3)**3)**(1/2))**(1/3) +\
                    o * (-q/2-((q/2)**2+(p/3)**3)**(1/2))**(1/3)
        case EquatorialWave.MixedRossbyGravity:
            omega = np.sqrt((k/2)**2+1) - k/2
        case EquatorialWave.Kelvin:
            omega = -k

    omega = omega.real
    return omega * np.sqrt(c*beta)

def wave_frequency(wave: EquatorialWave, k: Quantity |np.ndarray | float, n: int=1,
                     inverse_k: bool=True, H: Quantity|float=25, g: Quantity=None,
                     beta: Quantity=None) -> float:
    """
    根据波数计算赤道波动的频率，单位为 1/s=Hz。公式来自 Matsuno (1966)。
    如果波数 k 无单位，则认为是 '1/m'，如果想使用纬向波数，可以使用单位
    'zonal_wavenumber' 来定义。

    返回的频率与 wave_frequency 函数的频率相差 2 * PI 倍。

    >>> from moisten import unit
    >>> from moisten.equatorialWaves import EquatorialWave, wave_frequency
    >>> k = np.linspace(-20, 20, 100) * unit('zonal_wavenumber')
    >>> omega = wave_frequency(EquatorialWave.MixedRossbyGravity, k)
    >>> # 转为 1/day
    >>> omega = omega.to('1/day')

    Parameters
    ----------
    wave : EquatorialWave
        波动类型

        The type of wave. See the `EquatorialWave` enum for options.
    k : np.ndarray | float
        波数, 如果没有单位，则认为是 '1/m'。这里的波数是指单位距离内的波动周期次数，
        而不是相位差。

        The wavenumber, which is the number of waves per unit distance.
        If no unit is given, it is assumed to be in '1/m'.
    n : int, optional
        频散关系中的 n。当波动类型为 MRG 与 Kelvin 波时，该值无效。

        The value of n in the dispersion relation. This is ignored for MRG and Kelvin waves.
    inverse_k : bool, optional
        是否反转波数的正负方向，默认为 True。当为 True 时，输入的波数 > 0 表示波动
        向东传播，反之为 False 时，输入的波数 > 0 表示波动向西传播。

        Whether to invert the sign of the wavenumber. Default is True.
        When True, a positive wavenumber indicates eastward propagation,
        and when False, a positive wavenumber indicates westward propagation.
    H : Quantity | float, optional
        相当深度，默认单位为米。

        Equivalent depth, default unit is meters.
    g : Quantity, optional
        重力加速度，如果不提供，则使用默认的地球重力加速度。

        Gravitational acceleration. If not provided, the earth's gravitational
        acceleration is used.
    beta : Quantity, optional
        赤道上的罗斯贝参数，如果不提供，则使用默认的地球上的罗斯贝参数。

        Rossby parameter at the equator. If not provided, the default value
        for the earth is used.

    Returns
    -------
    Quantity
        波动的频率，单位为 1/s = Hz。是单位时间内的波动次数。

        The angular frequency of the wave, in radians per second.
    """
    return wave_omega(wave, k, n, inverse_k, H, g, beta) / (2*PI)


def equivalent_depth(wave: EquatorialWave,  k: Quantity | np.ndarray | float,
                     freq: Quantity | np.ndarray | float, n: int=1,
                     inverse_k: bool=True, g: Quantity=None,
                     beta: Quantity=None, depth_lim: Quantity=None) -> Quantity:
    """根据波数和频率计算赤道波动的相当深度，单位为米，便于画深度等值线。
    公式来自 Matsuno (1966)。

    >>> k = np.linspace(-20, 20, 100) * unit('zonal_wavenumber')
    >>> omega = np.linspace(0, 0.8, 100) * unit('1/day')
    >>> K, O = np.meshgrid(k, omega)
    >>> depth = equivalent_depth(EquatorialWave.MixedRossbyGravity, K, O)
    >>>
    >>> # 画出频率-波数图的等深度线
    >>> fig, ax = plt.subplots()
    >>> C = ax.contour(k, omega, depth, levels=[12, 25, 50, 90], colors='black')
    >>> ax.clabel(C, inline=True, fontsize=8)
    >>> plt.show()

    Calculate the equivalent depth of equatorial waves based on wavenumber and frequency.
    The formula is from Matsuno (1966). The equivalent depth is used to plot
    depth contours.

    Parameters
    ----------
    wave : EquatorialWave
        波动类型

        The type of wave. See the `EquatorialWave` enum for options.
    k : Quantity | np.ndarray | float
        波数, 如果没有单位，则认为是 '1/m'。这里的波数是指单位距离内的波动周期次数，
        而不是相位差。

        The wavenumber, which is the number of waves per unit distance.
    freq : Quantity | np.ndarray | float
        频率, 如果没有单位，则认为是 '1/s'。

        The frequency, which is the number of waves per unit time.
    n : int, optional
        频散关系中的 n，当波动类型为 MRG 与 Kelvin 波时，该值无效，默认为 1

        The value of n in the dispersion relation. This is ignored for MRG and Kelvin waves.
    inverse_k : bool, optional
        是否反转波数的正负方向，默认为 True

        Whether to invert the sign of the wavenumber. Default is True.
    g : Quantity, optional
        重力加速度，如果不提供，则使用默认的地球重力加速度。

        Gravitational acceleration. If not provided, the earth's gravitational
    beta : Quantity, optional
        赤道上的罗斯贝参数，如果不提供，则使用默认的地球上的罗斯贝参数。

        Rossby parameter at the equator. If not provided, the default value
        for the earth is used.
    depth_lim : Quantity, optional
        深度限制，会裁剪超过此深度的值，默认值为 1000 米。

        Depth limit for the equivalent depth. Values exceeding this depth
        will be clipped. Default is 1000 meters.

    Returns
    -------
    Quantity
        相当深度的计算结果，单位为米。
    """

    # 把 k 转为波数（个/米）
    k = _handle_quantity(k, '1/m')
    # 转为相位/米的波数
    k = k * 2 * PI
    freq = _handle_quantity(freq, '1/s')
    omega = freq * 2 * PI

    # 处理常数
    if g is None:
        g = EARTH_GRAVITY
    if beta is None:
        beta = ROSSBY_PARAMETER_ON_EQUATOR
    if depth_lim is None:
        depth_lim = 1000 * unit('m')

    if inverse_k:
        k = -k

    k = k

    match wave:
        case EquatorialWave.EastwardInertioGravity:
            k, omega = np.meshgrid(k, omega)
            h = (2*n+1)**2-4*omega**2/beta*(k/omega-k**2/beta)
            h = ((2*n+1)-h)/(2*(k/omega-k**2/beta))
            h = h**2/g
            return h
        case EquatorialWave.WestwardInertioGravity:
            if n == 0:
                raise ValueError("n=0 is not supported for WIG wave.")
            elif n > 0:
                return ((-beta*(2*n+1)+np.sqrt(beta**2 * (2*n+1)**2 + 4* k**2 * omega**2 ))/( 2*k**2))**2/g
        case EquatorialWave.EquatorialRossby:
            if n == 0:
                raise ValueError("n=0 is not supported for Rossby wave.")
            elif n > 0:
                h = ((omega*beta*(2*n+1))/(k*(beta-omega*k)))**2/g
                Omega = wave_omega(EquatorialWave.EquatorialRossby, k[0]/(2*PI), n, H=depth_lim,
                                   g=g, beta=beta, inverse_k=not inverse_k)
                h[omega > Omega[..., :]] = np.nan
                return h
        case EquatorialWave.MixedRossbyGravity:
            h = omega**4/((beta-omega*k)**2*g)
            Omega = wave_omega(EquatorialWave.MixedRossbyGravity, k[0]/(2*PI), n,
                                H=depth_lim, g=g, beta=beta, inverse_k=not inverse_k)
            h[omega > Omega[..., :]] = np.nan
            return h
        case EquatorialWave.Kelvin:
            h = omega**2/k**2/g
            Omega = wave_omega(EquatorialWave.Kelvin, k[0]/(2*PI), n,
                              H=depth_lim, g=g, beta=beta, inverse_k=not inverse_k)
            h[omega > Omega[..., :]] = np.nan
            return h
